import os
import multiprocessing as mp
from tqdm import tqdm as tqdm
from catalog.place import PlaceRecord, PlaceAPI
from common.geoimage.raster_dataset import RasterDataset, mosaic_raster_datasets
from analytics.crop_identifier_system.preprocessing.MosaicCropHelper import (
    MosaicCropHelper,
)
import analytics.crop_recognition.common_settings.model_config as config


def mosaic_crop_based_on_item(
    item: PlaceRecord, year, src_folder, dest_folder, level=2
):
    """
    Mosaic and crop to AoI boundary

    Parameters
    ----------
    item : PlaceRecord
        PlaceRecord item containing geom attribtue
    src_folder : str
        Path to folder which contains needed tiles with name as tileids.
    dest_folder : str
        Path to destination folder.
    """

    file_base_name = "_".join([item.parent_name_ori.replace(":", "_"), item.name_ori])
    if level == 2:
        file_base_name = file_base_name.replace("_不统计", "")
    dest_json_name = f"{file_base_name}.geojson"
    dest_raster_name = f"{file_base_name}.tif"

    individual_dest_folder = os.path.join(dest_folder)
    if not os.path.exists(individual_dest_folder):
        os.makedirs(individual_dest_folder, exist_ok=True)

    dest_json_path = os.path.join(individual_dest_folder, dest_json_name)
    dest_raster_path = os.path.join(individual_dest_folder, dest_raster_name)
    geom = item.geom
    mgrs_set = PlaceAPI.get_mgrs_set_by_geom(geom)
    mgrs_list = sorted(list(mgrs_set))
    path_list = []
    for mgrs_tileid in mgrs_list:
        file_path = os.path.join(src_folder, f"{mgrs_tileid}_{year}_class.tif")
        path_list.append(file_path)

    mosaic_rst = MosaicCropHelper.mosaic_crop_rasterdataset(
        path_list,
        geom,
        dest_res=10,
    )
    mask = mosaic_rst.data > 1
    # mosaic_rst.to_polygon(dest_json_path, mask_band=mask)
    tif_rst = mosaic_rst.paint(geom, 0, inverse=True)
    tif_rst.to_geotiff(dest_raster_path)


if __name__ == "__main__":
    year_list = [2023]  # [2018, 2019, 2020, 2021]
    items = []
    level = 2
    region = "辽宁"
    if level == 2:
        cur_items = PlaceRecord.query_many_items(
            target_parent_name=region, region_level=2
        )
        for item in PlaceRecord.query_many_items(
            target_parent_name=region, region_level=3
        ):
            if item.parent_name_ori.split(":")[2] == "不统计":
                cur_items.append(item)
    else:
        cur_items = PlaceRecord.query_many_items(
            target_parent_name=region, region_level=3
        )
    src_folder_root = f"/your_path_to_load_inference/{config.region}/{config.name}/processed/10"  # noqa:E501
    dest_folder_root = f"/your_path_to_load_inference/{config.region}/{config.name}/delivery_city"  # noqa; E501

    items.extend(cur_items)
    print(len(items), "items need to be processed")
    for year in year_list:
        year_str = str(year)
        year_int = int(year)
        src_folder = os.path.join(src_folder_root, year_str)
        dest_folder = os.path.join(dest_folder_root, year_str)
        if not os.path.exists(dest_folder):
            os.makedirs(dest_folder, exist_ok=True)
        pool = mp.Pool(30)
        pbar = tqdm(total=len(items))

        def update_pbar(*args):
            pbar.update(1)

        for item in items:
            pool.apply_async(
                mosaic_crop_based_on_item,
                args=[item, year, src_folder, dest_folder],
                callback=update_pbar,
                error_callback=print,
            )
        pool.close()
        pool.join()
